package com.openess.bigsearch.engine.index.merge;

import java.io.IOException;
import java.io.Writer;
import java.util.Arrays;
import java.util.Collection;
import java.util.Comparator;
import java.util.Iterator;

/**
 * 败者树，主要用于外排序和多路归并结合，可以提高排序效率，也用于主索引合并。
 * 
 * @author xiaoshuo
 * 
 */
public class LoserTree<T> {
	/**
	 * 败者树的非叶子节点，用于记录nodes中元素的下标，除了tree[0]是最终胜者外，其他位置都保存败者
	 */
	private int[] tree;
	/**
	 * 败者树的叶子节点，叶子节点和待排序的源数据段是一一对应的, 即第一个叶子节点记录第一个数据段的当前数据
	 */
	private Object[] nodes;
	/**
	 * 所有待排序源数据段，必须已排序，为nodes不断提供数据，位置必须和nodes对应
	 */
	private Iterator<T>[] branches;
	private Comparator<? super T> comparator;

	/**
	 * 构造方法，自定义比较器
	 * 
	 * @param comparator
	 *            排序数据的比较器
	 * @param branches
	 *            待排序的数据段，必须支持迭代，且数据段要求已排序
	 */
	public LoserTree(Comparator<? super T> comparator, Iterator<T>... branches) {
		this.comparator = comparator;
		this.branches = branches;
		init();
	}

	/**
	 * 构造方法，利用Comparable接口来排序
	 * 
	 * @param branches
	 */
	public LoserTree(Iterator<T>... branches) {
		this(new Comparator<T>() {
			@SuppressWarnings("unchecked")
			@Override
			public int compare(T o1, T o2) {
				return ((Comparable<T>) o1).compareTo(o2);
			}
		}, branches);
	}

	@SuppressWarnings("unchecked")
	public LoserTree(Collection<Iterator<T>> branches) {
		this(branches.toArray(new Iterator[0]));
	}

	@SuppressWarnings("unchecked")
	public LoserTree(Comparator<T> comparator, Collection<Iterator<T>> branches) {
		this(comparator, branches.toArray(new Iterator[0]));
	}

	/**
	 * 取出冠军数据，并重新调节败者树
	 */
	public T pop() {
		T result = get(tree[0]);
		if (result != null) {
			put(tree[0]);
			// 因为tree中的其他支路已经固定，因此仅需要比较变化的那一条支路
			// 因此除了初始化以外，每次pop仅需要O(lgN)次比较，重新计算tree[0]的值
			adjust(tree[0]);
		}
		return result;
	}

	/**
	 * 仅仅查看冠军数据
	 * 
	 * @return
	 */
	public T top() {
		return get(tree[0]);
	}

	/**
	 * 多路归并排序，并写入磁盘
	 * 
	 * @param out
	 * @throws IOException
	 */
	public void merge(Writer out) throws IOException {
		T top = null;
		while ((top = pop()) != null)
			out.write(top.toString() + " ");
	}

	/**
	 * 败者树的初始化操作
	 */
	private void init() {
		int size = branches.length;
		this.tree = new int[size];
		this.nodes = new Object[size];
		// 给所有叶子节点赋值
		for (int i = 0; i < size; i++)
			put(i);
		int winner = 0;
		// 计算出nodes中最小的下标
		for (int i = 1; i < size; i++)
			if (beat(i, winner))
				winner = i;
		// tree全部初始化为winner
		Arrays.fill(tree, winner);
		// 从后向前依次调整非叶子节点
		for (int i = size - 1; i >= 0; i--)
			adjust(i);
	}

	/**
	 * 调整第index个叶子节点<br>
	 * 具体调整过程为: 叶子节点和父节点比较, 败者留在父节点位置, 胜者继续和父节点的父节点比较,直到整棵树的根节点
	 * 
	 * @param index
	 */
	private void adjust(int index) {
		int size = branches.length;
		int t = (size + index) / 2;
		while (t > 0) {
			// 败者留在父节点的位置，tree中除了0以外，其他位置都是败者。
			// index代表胜者，每循环一次，就将子层的胜者和父层的败者进行一个比较，重新将index标记为胜者并在tree中标记败者
			if (beat(tree[t], index)) {
				int win = tree[t];
				tree[t] = index;
				index = win;
			}
			t = t / 2;
		}
		tree[0] = index;
	}

	/**
	 * 从指定的branches取数据，并放入对应的nodes中，如果指定的branches没有数据，那么nodes中插入null
	 * 
	 * @param index
	 *            指定的branches下标
	 */
	private void put(int index) {
		Iterator<T> branch = branches[index];
		nodes[index] = branch.hasNext() ? branch.next() : null;
	}

	/**
	 * 判断index1对应的节点是否能打败index2对应的节点，胜者为较小值，败者为较大值
	 * 
	 * @param index1
	 * @param index2
	 * @return
	 */
	private boolean beat(int index1, int index2) {
		T t1 = get(index1);
		T t2 = get(index2);
		if (t1 == null)
			return false;
		if (t2 == null)
			return true;
		// 这里, 当叶节点数据相等时比较分支索引是为了实现排序算法的稳定性
		int n = comparator.compare(t1, t2);
		return n != 0 ? n < 0 : index1 < index2;
	}

	/**
	 * 获取第index个叶子节点的当前数据<br>
	 * 如果返回null,则表示对应的数据源段没有数据了
	 * 
	 * @param index
	 * @return
	 */
	@SuppressWarnings("unchecked")
	private T get(int index) {
		return (T) nodes[index];
	}
}
